import io
import os
import uuid
from typing import Dict, Union, Optional

import requests
from fastapi import HTTPException, UploadFile, File
from fastapi.responses import JSONResponse, Response
from PIL import Image
from rembg import remove, new_session
from app.setting import settings

class KouTuAPI:
    def __init__(self):
        self.static_path = f"{settings.STATIC_PATH}/images"
        self.static_dir = f"{settings.STATIC_DIR}/images"
        # 确保静态资源目录存在
        os.makedirs(self.static_dir, exist_ok=True)

    def handle(self, url: str, type: str = "json", bg_color: str = None):
        """
        处理抠图请求
        
        Args:
            url: 远程图片URL
            type: 返回类型，'json'返回URL，'image'直接返回图片
            bg_color: 背景颜色，格式为十六进制颜色代码，如"#FF0000"
            
        Returns:
            根据type参数返回不同格式的结果
        """
        try:
            # 下载图片
            response = requests.get(url, timeout=10)
            if response.status_code != 200:
                # raise HTTPException(status_code=400, detail="无法下载图片")
                return JSONResponse({
                    "success": False,
                    "message": "无法下载图片",
                }, status_code=400, media_type="application/json, charset=utf-8")
            
            # 处理图片
            input_image = Image.open(io.BytesIO(response.content))
            output_image = self.remove_background(input_image, bg_color)
            
            # 根据type返回不同格式
            if type.lower() == "json":
                # 保存图片并返回JSON
                image_path, image_url = self._save_image(output_image)
                return JSONResponse({
                    "success": True,
                    "data": {
                        "url": image_url
                    },
                    "code": 200
                }, status_code=200, media_type="application/json, charset=utf-8")
            elif type.lower() == "image":
                # 直接返回图片
                img_byte_arr = io.BytesIO()
                output_image.save(img_byte_arr, format='PNG')
                return Response(
                    content=img_byte_arr.getvalue(),
                    media_type="image/png",
                    headers={
                        "Content-Disposition": "attachment; filename=removed_bg.png"
                    },
                    status_code=200
                )
            else:
                # raise HTTPException(status_code=400, detail="不支持的返回类型")
                return JSONResponse({
                    "success": False,
                    "message": "不支持的返回类型",
                    "code": 400
                }, status_code=400, media_type="application/json, charset=utf-8")
                
        except Exception as e:
            # raise HTTPException(status_code=500, detail=f"处理图片时出错: {str(e)}")
            return JSONResponse({
                "success": False,
                "message": f"处理图片时出错: {str(e)}",
                "code": 500
            }, status_code=500, media_type="application/json, charset=utf-8")
    
    async def handle_upload(self, file: UploadFile = File(...), type: str = "json", bg_color: str = None):
        """
        处理用户上传的图片进行抠图
        
        Args:
            file: 用户上传的图片文件
            type: 返回类型，'json'返回URL，'image'直接返回图片
            bg_color: 背景颜色，格式为十六进制颜色代码，如"#FF0000"
            
        Returns:
            根据type参数返回不同格式的结果
        """
        try:
            # 读取上传的图片
            contents = await file.read()
            
            # 处理图片
            input_image = Image.open(io.BytesIO(contents))
            output_image = self.remove_background(input_image, bg_color)
            
            # 根据type返回不同格式
            if type.lower() == "json":
                # 保存图片并返回JSON
                image_path, image_url = self._save_image(output_image)
                return JSONResponse({
                    "success": True,
                    "data": {
                        "url": image_url
                    },
                    "code": 200
                })
            elif type.lower() == "image":
                # 直接返回图片
                img_byte_arr = io.BytesIO()
                output_image.save(img_byte_arr, format='PNG')
                img_byte_arr.seek(0)
                return Response(
                    content=img_byte_arr.getvalue(),
                    media_type="image/png",
                    headers={
                        "Content-Disposition": "attachment; filename=removed_bg.png"
                    },
                    status_code=200
                )
            else:
                # raise HTTPException(status_code=400, detail="不支持的返回类型")
                return JSONResponse({
                    "success": False,
                    "message": "不支持的返回类型",
                    "code": 400
                }, status_code=400, media_type="application/json, charset=utf-8")
                
        except Exception as e:
            # raise HTTPException(status_code=500, detail=f"处理图片时出错: {str(e)}")
            return JSONResponse({
                "success": False,
                "message": f"处理图片时出错: {str(e)}",
                "code": 500
            }, status_code=500, media_type="application/json, charset=utf-8")
    
    def remove_background(self, image: Image.Image, bg_color: str = None) -> Image.Image:
        """
        移除图片背景
        
        Args:
            image: PIL图像对象
            bg_color: 背景颜色，格式为十六进制颜色代码，如"#FF0000"
            
        Returns:
            处理后的PIL图像对象
        """
        # u2net, u2netp, u2net_human_seg, silueta
        # 各模型特点：
        # u2net：通用模型，质量高但速度较慢（默认）
        # u2netp：轻量版，速度较快
        # u2net_human_seg：专门用于人像
        # silueta：最小模型，速度最快
        # 选择模型
        session = new_session(model_name="u2net")
        
        # 处理背景色
        bgcolor = None
        if bg_color:
            try:
                # 移除可能的#前缀
                if bg_color.startswith('#'):
                    bg_color = bg_color[1:]
                    
                # 解析颜色
                r = int(bg_color[0:2], 16)
                g = int(bg_color[2:4], 16)
                b = int(bg_color[4:6], 16)
                
                bgcolor = (r, g, b, 255)
            except Exception as e:
                bgcolor = None
        
        return remove(image, session=session, bgcolor=bgcolor)
    
    def _save_image(self, image: Image.Image) -> tuple:
        """
        保存图片并返回路径和URL
        
        Args:
            image: PIL图像对象
            
        Returns:
            (文件路径, 图片URL)的元组
        """
        # 生成唯一文件名
        filename = f"{uuid.uuid4()}.png"
        file_path = os.path.join(self.static_dir, filename)
        
        # 保存图片
        image.save(file_path, "PNG")
        
        # 生成URL
        image_url = f"{self.static_path}/{filename}"
        
        return file_path, image_url
 